{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zqc2vWQXPzMx"
      },
      "source": [
        "This notebook generates sequence tagging examples for the COGS dataset."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vuTAjNNzMblE"
      },
      "outputs": [],
      "source": [
        "# Licensed under the Apache License, Version 2.0\n",
        "\n",
        "import collections\n",
        "import copy\n",
        "import dataclasses\n",
        "import json\n",
        "from typing import List, Mapping, Sequence, Tuple\n",
        "\n",
        "import tensorflow as tf  # Only needed for GFile."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ev3oQIwaMfXp"
      },
      "outputs": [],
      "source": [
        "DATA_DIR = ''  # Directory containing original tsv data files.\n",
        "OUTPUT_DIR = ''  # Directory to write sequence tagged json files.\n",
        "DATA_FILES = [f'{x}.tsv' for x in ('dev', 'gen', 'test', 'train', 'train_100')]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qG0CC1NKWaZG"
      },
      "outputs": [],
      "source": [
        "@dataclasses.dataclass\n",
        "class RawExample:\n",
        "  input: str\n",
        "  target: str\n",
        "  distribution: str\n",
        "\n",
        "@dataclasses.dataclass\n",
        "class LabeledExample:\n",
        "  tokens: List[str]  # Raw input tokens.\n",
        "  parent: List[int]  # Index of parent token or `-1` for no parent.\n",
        "  role: List[str]  # Role of token.\n",
        "  category: List[str]  # Only \"VERB\" for verbs and \"CNOUN\" for common nouns; other words get empty category.\n",
        "  noun_type: List[str]  # Whether this is a definite or indefinite noun (or neither).\n",
        "  verb_name: List[str]  # Name of verb if token is a verb.\n",
        "  distribution: str"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fbuBxs8rQwrV"
      },
      "outputs": [],
      "source": [
        "def tokenize_input(input: str) -\u003e List[str]:\n",
        "  return input.split(' ')\n",
        "\n",
        "def split_clauses(target: str) -\u003e List[str]:\n",
        "  compact_target = ''.join(target.split(' '))\n",
        "  targets = compact_target.split('AND')\n",
        "  targets = targets[0].split(';') + targets[1:]\n",
        "  return targets\n",
        "\n",
        "def get_arguments(clause: str) -\u003e List[str]:\n",
        "  assert clause[-1] == ')'\n",
        "  idx = clause.find('(')\n",
        "  assert idx != -1\n",
        "  args = clause[(idx + 1):-1].split(',')\n",
        "  assert 1 \u003c= len(args) \u003c= 2\n",
        "  return args\n",
        "\n",
        "def get_predicate(clause: str) -\u003e str:\n",
        "  start_idx = 1 if clause.startswith('*') else 0\n",
        "  end_idx = clause.find('(')\n",
        "  assert end_idx != -1\n",
        "  return clause[start_idx:end_idx]\n",
        "\n",
        "def is_index_arg(arg: str) -\u003e bool:\n",
        "  return arg.startswith('x_')\n",
        "\n",
        "def get_arg_index(arg: str) -\u003e int:\n",
        "  return int(arg[2:])\n",
        "\n",
        "def is_noun_clause(clause: str) -\u003e bool:\n",
        "  return '.' not in clause\n",
        "\n",
        "def get_preposition(clause: str) -\u003e str:\n",
        "  marker = '.nmod.'\n",
        "  idx = clause.find(marker)\n",
        "  assert idx != -1\n",
        "  return clause.split('(')[0][idx + len(marker):]\n",
        "\n",
        "def is_proper_noun(token: str) -\u003e bool:\n",
        "  return token[0].isupper() and token not in ('A', 'An', 'The', 'TV')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oUR976j6ch-D"
      },
      "outputs": [],
      "source": [
        "NONE_ROLE = ''\n",
        "\n",
        "class LabeledExampleBuilder:\n",
        "\n",
        "  def __init__(self, raw_example: RawExample):\n",
        "    tokens = tokenize_input(raw_example.input)\n",
        "    self._example = LabeledExample(\n",
        "        tokens=list(tokens),\n",
        "        parent=[-1] * len(tokens),\n",
        "        role=[NONE_ROLE] * len(tokens),\n",
        "        category=[NONE_ROLE] * len(tokens),\n",
        "        noun_type=[NONE_ROLE] * len(tokens),\n",
        "        verb_name=[NONE_ROLE] * len(tokens),\n",
        "        distribution=raw_example.distribution)\n",
        "\n",
        "    self._idx_by_token = collections.defaultdict(list)\n",
        "    for i, token in enumerate(tokens):\n",
        "      self._idx_by_token[token].append(i)\n",
        "\n",
        "    if self._process_primitive_example(raw_example):\n",
        "      return\n",
        "\n",
        "    self._mark_proper_noun_category()\n",
        "    clauses = split_clauses(raw_example.target)\n",
        "    clauses = self._remove_redundant_xcomp_agent(clauses)\n",
        "    clauses = self._process_nouns(clauses)\n",
        "    clauses = self._process_prepositions(clauses)\n",
        "    # All remaining clauses should be for verbs now.\n",
        "    self._process_verbs(clauses)\n",
        "\n",
        "  def build(self) -\u003e LabeledExample:\n",
        "    return copy.deepcopy(self._example)\n",
        "\n",
        "  def _process_primitive_example(self, raw_example: RawExample) -\u003e bool:\n",
        "    \"\"\"Special handling if `raw_example` is primitive; no-op otherwise.\n",
        "\n",
        "    Primitive examples are just single-word inputs; either a single noun\n",
        "    or single verb in particular.  They can have LAMBDA notation in the target\n",
        "    that we'll basically ignore for sequence labeling.\n",
        "\n",
        "    Here's an example of a primitive example:\n",
        "      input: touch\n",
        "      target: LAMBDA a . LAMBDA b . LAMBDA e . touch . agent ( e , b ) AND touch . theme ( e , a )\n",
        "      distribution: primitive\n",
        "\n",
        "    Returns True if `raw_example` is primitive; False otherwise.\n",
        "    \"\"\"\n",
        "    if raw_example.distribution != 'primitive':\n",
        "      return False\n",
        "\n",
        "    assert len(self._example.tokens) == 1\n",
        "\n",
        "    marker_str = 'LAMBDA'\n",
        "    idx1 = raw_example.target.find(marker_str)\n",
        "    if idx1 == -1:\n",
        "      # Proper nouns don't have any \"LAMBDA\"\n",
        "      self._example.category[0] = 'PNOUN'\n",
        "      return True\n",
        "    idx2 = raw_example.target.find(marker_str, idx1 + 1)\n",
        "\n",
        "    if idx2 == -1:\n",
        "      # Only 1 \"LAMBDA\", so the primitive is a common noun.\n",
        "      self._example.category[0] = 'CNOUN'\n",
        "    else:\n",
        "      # More than 1 \"LAMBDA\", so the primitive is a verb.\n",
        "      self._example.category[0] = 'VERB'\n",
        "      verb = self._example.tokens[0]\n",
        "      assert f' {verb} ' in raw_example.target\n",
        "      self._example.verb_name[0] = verb\n",
        "\n",
        "    return True\n",
        "\n",
        "  def _mark_proper_noun_category(self) -\u003e None:\n",
        "    for i, token in enumerate(self._example.tokens):\n",
        "      if is_proper_noun(token):\n",
        "        self._example.category[i] = 'PNOUN'\n",
        "\n",
        "  def _remove_redundant_xcomp_agent(self, clauses: List[str]) -\u003e List[str]:\n",
        "    \"\"\"Removes redundant agent clauses for xcomp relations.\n",
        "\n",
        "    For example, for input \"Audrey wished to crawl .\", the clauses are\n",
        "    [\"wish.agent(x_1,Audrey)\", \"wish.xcomp(x_1,x_3)\", \"crawl.agent(x_3,Audrey)\"]\n",
        "\n",
        "    We remove \"crawl.agent(x_3,Audrey)\" since it's already implied by the\n",
        "    other two clauses.  This is necessary since \"Audrey\" would otherwise\n",
        "    have both \"wish\" and \"crawl\" as parents.  We only want \"wish\" as the\n",
        "    parent.\n",
        "    \"\"\"\n",
        "    xcomp_clauses = [x for x in clauses if '.xcomp' in x]\n",
        "    for clause in xcomp_clauses:\n",
        "      xcomp_args = get_arguments(clause)\n",
        "      assert len(xcomp_args) == 2\n",
        "      agent_clause_prefix = clause.split(',')[0].replace('xcomp', 'agent')\n",
        "      agent_clauses = [x for x in clauses if x.startswith(agent_clause_prefix)]\n",
        "      assert len(agent_clauses) == 1\n",
        "      agent_args = get_arguments(agent_clauses[0])\n",
        "      assert len(agent_args) == 2\n",
        "\n",
        "      # Construct the clause we wish to remove.\n",
        "      infinitive_verb = self._example.tokens[get_arg_index(xcomp_args[-1])]\n",
        "      redundant_clause = (\n",
        "          f'{infinitive_verb}.agent({xcomp_args[-1]},{agent_args[-1]})')\n",
        "      clauses.remove(redundant_clause)\n",
        "\n",
        "      # Set the verb_name for the infinitive verb.\n",
        "      infinitive_idx = get_arg_index(xcomp_args[-1])\n",
        "      self._example.verb_name[infinitive_idx] = infinitive_verb\n",
        "    return clauses\n",
        "\n",
        "  def _process_nouns(self, clauses: List[str]) -\u003e List[str]:\n",
        "    \"\"\"Processes all noun clauses and returns all remaining clauses.\n",
        "\n",
        "    This only populates `noun_type` labels.\n",
        "    \"\"\"\n",
        "    remaining_clauses = []\n",
        "    for clause in clauses:\n",
        "      if is_noun_clause(clause):\n",
        "        args = get_arguments(clause)\n",
        "        assert len(args) == 1\n",
        "        assert is_index_arg(args[0])\n",
        "        idx = get_arg_index(args[0])\n",
        "        self._example.category[idx] = 'CNOUN'\n",
        "        assert self._example.noun_type[idx] == NONE_ROLE\n",
        "        if clause.startswith('*'):\n",
        "          self._example.noun_type[idx] = 'DEF'\n",
        "        else:\n",
        "          self._example.noun_type[idx] = 'INDEF'\n",
        "      else:\n",
        "        remaining_clauses.append(clause)\n",
        "\n",
        "    # Sanity check that all remaining clauses have more than one argument.\n",
        "    for clause in remaining_clauses:\n",
        "      assert ',' in clause\n",
        "\n",
        "    return remaining_clauses\n",
        "\n",
        "  def _process_prepositions(self, clauses: List[str]) -\u003e List[str]:\n",
        "    remaining_clauses = []\n",
        "    for clause in clauses:\n",
        "      if '.nmod.' in clause:\n",
        "        args = get_arguments(clause)\n",
        "        assert len(args) == 2\n",
        "        parent_noun_idx = get_arg_index(args[0])\n",
        "        pp_noun_idx = get_arg_index(args[1])  # Prepositional phrase noun.\n",
        "\n",
        "        # Find index of preposition\n",
        "        preposition = get_preposition(clause)\n",
        "        candidate_indices = self._idx_by_token[preposition]\n",
        "        indices = [x for x in candidate_indices\n",
        "                   if parent_noun_idx \u003c x \u003c pp_noun_idx]\n",
        "        assert len(indices) == 1\n",
        "        prep_idx = indices[0]\n",
        "\n",
        "        # Make the preposition the parent of the prepositional phrase noun\n",
        "        # and the parent noun the parent of the preposition.\n",
        "        assert self._token_is_unassigned(pp_noun_idx)\n",
        "        self._assign_token(pp_noun_idx, 'PP_NOUN', prep_idx)\n",
        "        assert self._token_is_unassigned(prep_idx)\n",
        "        self._assign_token(prep_idx, 'PREP', parent_noun_idx)\n",
        "      else:\n",
        "        remaining_clauses.append(clause)\n",
        "    return remaining_clauses\n",
        "\n",
        "  def _process_verbs(self, clauses: List[str]) -\u003e None:\n",
        "    remaining_clauses = []\n",
        "    for clause in clauses:\n",
        "      args = get_arguments(clause)\n",
        "      assert len(args) == 2\n",
        "      verb_idx = get_arg_index(args[0])\n",
        "      if is_index_arg(args[1]):\n",
        "        child_idx = get_arg_index(args[1])\n",
        "      else:  # Proper noun argument.\n",
        "        indices = self._idx_by_token[args[1]]\n",
        "        assert len(indices) == 1\n",
        "        child_idx = indices[0]\n",
        "      predicate = clause.split('(')[0]\n",
        "      verb, role = predicate.split('.')\n",
        "\n",
        "      assert self._token_is_unassigned(child_idx)\n",
        "      self._assign_token(child_idx, role, verb_idx)\n",
        "\n",
        "      curr_name = self._example.verb_name[verb_idx]\n",
        "      assert curr_name == NONE_ROLE or curr_name == verb\n",
        "      self._example.verb_name[verb_idx] = verb\n",
        "      self._example.category[verb_idx] = 'VERB'\n",
        "\n",
        "  def _token_is_unassigned(self, index: int) -\u003e bool:\n",
        "    return (self._example.role[index] == NONE_ROLE and\n",
        "            self._example.parent[index] == -1)\n",
        "\n",
        "  def _assign_token(self, index: int, role: str, parent: int) -\u003e None:\n",
        "    self._example.role[index] = role\n",
        "    self._example.parent[index] = parent"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rmh929JvBrm5"
      },
      "outputs": [],
      "source": [
        "# Utilities to confirm that we can reconstruct the target from sequence labels.\n",
        "\n",
        "@dataclasses.dataclass\n",
        "class WrappedArg:\n",
        "  index: int\n",
        "  name: str = ''  # Only present for proper nouns.\n",
        "\n",
        "  def to_str(self):\n",
        "    if self.name:\n",
        "      return self.name\n",
        "    return f'x _ {self.index}'\n",
        "\n",
        "@dataclasses.dataclass\n",
        "class ReconstructedClause:\n",
        "  predicate: str\n",
        "  arguments: List[WrappedArg]\n",
        "  is_definite: bool = False\n",
        "\n",
        "  def to_str(self):\n",
        "    arg_str = ' , '.join(x.to_str() for x in self.arguments)\n",
        "    result = f'{self.predicate} ( {arg_str} )'\n",
        "    if self.is_definite:\n",
        "      result = '* ' + result\n",
        "    return result\n",
        "\n",
        "\n",
        "def reconstruct_target(example: LabeledExample) -\u003e str:\n",
        "  # Create noun clauses first, handling definite nouns differently\n",
        "  # since they're separated and joined with ';' instead of 'AND'.\n",
        "  definite_clauses = []\n",
        "  other_clauses = []\n",
        "  for i, noun_type in enumerate(example.noun_type):\n",
        "    predicate = example.tokens[i]\n",
        "    if predicate == 'TV':\n",
        "      predicate = 'tv'\n",
        "    if noun_type == 'DEF':\n",
        "      definite_clauses.append(\n",
        "          ReconstructedClause(predicate, [WrappedArg(i)], True))\n",
        "    elif noun_type == 'INDEF':\n",
        "      other_clauses.append(\n",
        "          ReconstructedClause(predicate, [WrappedArg(i)]))\n",
        "\n",
        "  # Add all other clauses, which correspond to non-empty roles.\n",
        "  for i, role in enumerate(example.role):\n",
        "    if role == 'PREP':\n",
        "      # We capture the whole prepositional phrase clause from the 'PP_NOUN'\n",
        "      # so we skip 'PREP' tokens for now.\n",
        "      continue\n",
        "    elif role == 'PP_NOUN':\n",
        "      assert example.noun_type[i] != NONE_ROLE  # Never a proper noun.\n",
        "      prep_idx = example.parent[i]\n",
        "      preposition = example.tokens[prep_idx]\n",
        "      parent_noun_idx = example.parent[prep_idx]\n",
        "      parent_noun = example.tokens[parent_noun_idx]\n",
        "      if parent_noun == 'TV':\n",
        "        parent_noun = 'tv'\n",
        "      predicate = f'{parent_noun} . nmod . {preposition}'\n",
        "      assert example.noun_type[parent_noun_idx] != NONE_ROLE  # Never a proper noun.\n",
        "      other_clauses.append(\n",
        "          ReconstructedClause(\n",
        "              predicate, [WrappedArg(parent_noun_idx), WrappedArg(i)]))\n",
        "    elif role != NONE_ROLE:\n",
        "      # Reconstruct verb clause\n",
        "      verb_idx = example.parent[i]\n",
        "      verb = example.verb_name[verb_idx]\n",
        "      predicate = f'{verb} . {role}'\n",
        "      other_clauses.append(\n",
        "          ReconstructedClause(\n",
        "              predicate, [WrappedArg(verb_idx), wrap_arg(i, example)]))\n",
        "    if role == 'xcomp':\n",
        "      verb = example.verb_name[i]\n",
        "      parent_idx = example.parent[i]\n",
        "      agent_idx = -1\n",
        "      for j in range(len(example.tokens)):\n",
        "        if example.role[j] == 'agent' and example.parent[j] == parent_idx:\n",
        "          agent_idx = j\n",
        "      assert agent_idx != -1\n",
        "      other_clauses.append(\n",
        "          ReconstructedClause(\n",
        "              f'{verb} . agent', [WrappedArg(i), wrap_arg(agent_idx, example)]))\n",
        "\n",
        "  def sort_key(clause: ReconstructedClause) -\u003e Tuple[int, int]:\n",
        "    if len(clause.arguments) \u003e 1:\n",
        "      second_idx = clause.arguments[1].index\n",
        "    else:\n",
        "      second_idx = -1\n",
        "    return (clause.arguments[0].index, second_idx)\n",
        "\n",
        "  other_clauses.sort(key=sort_key)\n",
        "\n",
        "  clause_strings = [x.to_str() for x in definite_clauses]\n",
        "  clause_strings.append(' AND '.join([x.to_str() for x in other_clauses]))\n",
        "  return ' ; '.join(clause_strings)\n",
        "\n",
        "\n",
        "def wrap_arg(index: int, example: LabeledExample) -\u003e WrappedArg:\n",
        "  \"\"\"Creates a `WrappedArg` for the index, handling proper nouns appropriately.\n",
        "  \"\"\"\n",
        "  if example.tokens[index][0].isupper():\n",
        "    name = example.tokens[index]\n",
        "  else:\n",
        "    name = ''\n",
        "  return WrappedArg(index, name)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "R4mwZYh3MuX8"
      },
      "outputs": [],
      "source": [
        "# Label and write out all examples.\n",
        "\n",
        "%%time\n",
        "\n",
        "for filename in DATA_FILES:\n",
        "  raw_examples = []\n",
        "  with tf.io.gfile.GFile(f'{DATA_DIR}{filename}') as f:\n",
        "    for line in f.readlines():\n",
        "      fields = [x.strip() for x in line.split('\\t')]\n",
        "      raw_examples.append(RawExample(*fields))\n",
        "\n",
        "  examples = []\n",
        "  for raw_example in raw_examples:\n",
        "    example = LabeledExampleBuilder(raw_example).build()\n",
        "    if raw_example.distribution != 'primitive':\n",
        "      assert reconstruct_target(example) == raw_example.target\n",
        "    examples.append(example)\n",
        "\n",
        "  out_filename = filename.split('.')[0] + '_seqtag.jsonl'\n",
        "  out_path = f'{OUTPUT_DIR}{out_filename}'\n",
        "  with tf.io.gfile.GFile(out_path, mode='w') as f:\n",
        "    f.write('\\n'.join(json.dumps(x.__dict__) for x in examples))"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "last_runtime": {
        "build_target": "//learning/deepmind/public/tools/ml_python:ml_notebook",
        "kind": "private"
      },
      "name": "cogs_seq_tagging_data_gen.ipynb",
      "private_outputs": true,
      "provenance": [
        {
          "file_id": "12EoeMlPWDOqjkm6-Uy_bFZiDyRcLXOfH",
          "timestamp": 1631135365807
        },
        {
          "file_id": "1yXH6yBFK8to69QyFTpe911-EiVL16peE",
          "timestamp": 1630515720092
        }
      ]
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
